dataset1_filtered <- dataset1_cleaned %>%
  filter(SHRCD == 73 | SHRCD == 11) %>% 
  filter(!is.na(RET) & !is.na(PRC))

cat("Number of unique security in the dataset: ",length(unique(dataset1_filtered$PERMNO)))

dataset1_df1 <- dataset1_filtered %>%
  arrange(PERMNO, date) %>%
  mutate(
    # Compute the dollar volume
    dollar_vol = PRC * VOL,
    
    # Compute the bid-ask spread (relative)
    bid_ask_spread = (ASK - BID) / ((ASK + BID) / 2),
    
    # Compute the turnover ratio
    turnover_ratio = VOL / (SHROUT * 1000)
  )

# Add Risk free rate (US Daily Treasury Bill Rates)
dataset1_df2 <- dataset1_df1 %>%
  arrange(date) %>%
  left_join(
    risk_free_rate %>% 
      filter(
        date >= "2019-08-20" & date <= "2020-08-20"
      )
    , by = "date"
  ) %>%
  mutate(
    risk_free_rate = na.locf(risk_free_rate, na.rm = FALSE)
  )

# Define the total trading days
trading_days_total <- length(unique(dataset1$date))

dataset1_df3 <- dataset1_df2 %>%
  group_by(PERMNO) %>%
  arrange(date) %>%
  mutate(
    # Compute daily mean return and standard deviation for each PERMNO
    mean_ret = rollmean(RET, k = 30, fill = NA, align = "right", na.rm = TRUE),
    std_ret = rollapply(RET, width = 30, FUN = sd, fill = NA, align = "right", na.rm = TRUE),
    
    # Compute volatility
    volatility = std_ret * sqrt(trading_days_total),
    
    # Compute daily Sharpe Ratio
    sharpe_ratio = ifelse(is.na(std_ret) | std_ret == 0, NA, (mean_ret - risk_free_rate) / std_ret)
    ) %>%
  ungroup() %>% 
  select(-c("mean_ret", "std_ret")) 

# Add sector ID
dataset1_df4 <- dataset1_df3 %>%
  left_join(industry_code, by = "TICKER") %>%
  mutate(
    SECTOR = case_when(
      SECTOR == "Health Care" ~ "Health_Care",
      SECTOR == "Information Technology" ~ "Information_Technology",
      SECTOR == "Consumer Discretionary" ~ "Consumer_Discretionary",
      SECTOR == "Consumer Staples" ~ "Consumer_Staples",
      SECTOR == "Real Estate" ~ "Real_Estate",
      SECTOR == "Communication Services" ~ "Communication_Services",
      TRUE ~ SECTOR  
    ),
    # Assign all missing value as Unknown Sector to avoid dropping large amount of data
    SECTOR = ifelse(is.na(SECTOR), "Unknown", SECTOR)
  ) %>% 
  select(-c(SECTOR_ID, INDUSTRY))

# Compute the sum of user holding of all securities
user_holding_summary <- userholding %>%
  group_by(timestamp) %>%
  summarise(total_holding = sum(users_holding, na.rm = TRUE), .groups = 'drop')

# Add total user holdings to the original dataset
user_holding_1 <- userholding %>%
  left_join(user_holding_summary, by = "timestamp")

# Compute market share
user_holding_2 <- user_holding_1 %>%
  group_by(TICKER, timestamp) %>% 
  mutate(
    market_share = users_holding / total_holding * 100) 

# Add the market_share to the dataset
dataset1_df5 <- dataset1_df4 %>%
  left_join(user_holding_2, by = c("TICKER" = "TICKER", "date" = "timestamp")) %>% 
  select(-c("users_holding", "total_holding"))

dataset1_df6 <- dataset1_df5 %>%
  filter(!is.na(SHROUT)) %>% 
   # Compute market capitalisation
  mutate(
    market_cap = SHROUT * PRC * 1000,
    #  Categorise company size for the security based on market capitalisation
    company_size = case_when(
      market_cap <= 2e9 ~ "Small",
      market_cap > 2e9 & market_cap <= 1e10 ~ "Medium",
      market_cap > 1e10 ~ "Large"
    )
  )

# Inspect the unique security count of each company size
print(dataset1_df6 %>%
  group_by(company_size) %>%
  summarise(distinct_tickers = n_distinct(PERMNO)))

dataset1_df7 <- dataset1_df6 %>%
  group_by(PERMNO) %>%
  mutate(
    abs_corr_sp = abs(cor(RET, sprtrn, use = "complete.obs")))

dataset1_df8 <- dataset1_df7 %>%
  mutate(security_type = case_when(
      SHRCD == 73 ~ "ETF",
      SHRCD == 11 ~ "Stock")
      )

# Tracking error
dataset1_df9 <- dataset1_df8 %>% 
  group_by(PERMNO) %>% 
  mutate(
    tracking_error = ifelse(is.na(sprtrn), NA, sd(RET - sprtrn))
  )

# Overall
write.csv(task1_overall_noncovid, file = "task1_overall_noncovid.csv", row.names = FALSE)
write.csv(task1_overall_covid, file = "task1_overall_covid.csv", row.names = FALSE)

# By Sector
write.csv(task1_sector_noncovid, file = "task1_sector_noncovid.csv", row.names = FALSE)
write.csv(task1_sector_covid, file = "task1_sector_covid.csv", row.names = FALSE)

# By Company size
write.csv(task1_size_noncovid, file = "task1_size_noncovid.csv", row.names = FALSE)
write.csv(task1_size_covid, file = "task1_size_covid.csv", row.names = FALSE)

# By security type
write.csv(task1_type_noncovid, file = "task1_type_noncovid.csv", row.names = FALSE)
write.csv(task1_type_covid, file = "task1_type_covid.csv", row.names = FALSE)

# Dataset for this plot is time series
time_series_plot_overall <- function(data, variables, var_names, periods) {
  
  # Convert the dataset to long format
  data_long <- data %>%
    pivot_longer(cols = all_of(variables), 
                 names_to = "variable", 
                 values_to = "value") %>%
    mutate(variable = factor(variable, 
                             levels = variables, 
                             labels = var_names)
           )
  
  # Create shaded areas for showing comparison periods
  shaded_areas <- bind_rows(
    lapply(periods, function(p) {
      data.frame(xmin = as.Date(p$start), 
                 xmax = as.Date(p$end), 
                 ymin = -Inf, 
                 ymax = Inf, 
                 fill = p$color, 
                 label = p$label)
    })
  )
  
  # Create plots for each variable
  create_plot <- function(var_name, dataset_name) {
    plot_data <- data_long %>% 
      filter(variable == var_name)
    
    p <- ggplot(plot_data, aes(x = date, 
                               y = value)) +
      
      geom_rect(data = shaded_areas, 
                aes(xmin = xmin, 
                    xmax = xmax, 
                    ymin = ymin, 
                    ymax = ymax, 
                    fill = label),
                alpha = 0.2, 
                inherit.aes = FALSE) +
      
      geom_line(aes(color = "Data"), 
                size = 1) +
      
      scale_fill_manual(values = setNames(unique(shaded_areas$fill),
                                          unique(shaded_areas$label)),
                        name = "Period") +
      
      guides(fill = guide_legend(title = "Periods for Comparison"),
             color = guide_legend(title = "Metrics")) +
      
      labs(title = paste("Overall Performance of", var_name, "during 20 August 2019 to 20 August 2020"),
           x = "Date", 
           y = paste("Standardised", var_name)) +
      
      scale_x_date(date_labels = "%Y-%m-%d", date_breaks = "1 month") +
      
      theme_minimal() +
      
      theme(axis.text.x = element_text(angle = 55, hjust = 1)) +
      
      scale_y_continuous(expand = expansion(mult = c(0.05, 0.1))) +
      
      scale_color_manual(name = "Metrics",
                         values = c("Data" = "red"),
                         labels = c("Data" = var_name))

    # Save plot for report
    save_plot(p, paste0(dataset_name, "_", var_name))
    
    return(p)
  }
  
  # To keep the saved plot file name easy to read
  dataset_name <- deparse(substitute(data))
  
  # Create a list of plots
  plots <- lapply(var_names, function(var_name) {
    create_plot(var_name, dataset_name)
  })
  
  # Print plots
  for (plot in plots) {
    print(plot)
  }
}

time_series_plot_group <- function(data, variables, var_names, label_name, periods, group, color_map) {
  
  # Convert the data to long format
  data_long <- data %>%
    pivot_longer(cols = all_of(variables), 
                 names_to = "variable", 
                 values_to = "value") %>%
    mutate(variable = factor(variable, 
                             levels = variables, 
                             labels = var_names)) %>%
    filter(!is.na(!!sym(group)))
  
  # Define color palette for the group variable
  group_colors <- brewer.pal(n = length(unique(data_long[[group]])), 
                             name = color_map)
  names(group_colors) <- unique(data_long[[group]])
  
  # Create shaded areas for showing comparison periods
  shaded_areas <- do.call(rbind, 
                          lapply(periods, function(p) {
                            data.frame(xmin = as.Date(p$start),
                                       xmax = as.Date(p$end), 
                                       ymin = -Inf, 
                                       ymax = Inf, 
                                       fill = p$color, 
                                       label = p$label)
                          })
  )
  
  # Create plots for each variable
  create_plot <- function(var_name, dataset_name) {
    plot_data <- data_long %>% 
      filter(variable == var_name)
    
    # Initialise the plot
    p <- ggplot() +
      
      # Loop over each group and plot its points and lines in order
      lapply(unique(plot_data[[group]]), function(g) {
        list(
          geom_point(data = plot_data %>% filter(!!sym(group) == g),
                     aes(x = date, y = value, color = !!sym(group)), size = 2, alpha = 0.9),
          geom_line(data = plot_data %>% filter(!!sym(group) == g),
                    aes(x = date, y = value, color = !!sym(group)), size = 1)
        )
      }) +
      
      # Create shaded comparison periods
      geom_rect(data = shaded_areas, 
                aes(xmin = xmin, 
                    xmax = xmax, 
                    ymin = ymin, 
                    ymax = ymax, 
                    fill = label),
                alpha = 0.2, 
                inherit.aes = FALSE) +
      
      # Customise colors and legends
      scale_color_manual(values = group_colors, 
                         name = label_name) +
      scale_fill_manual(values = setNames(unique(shaded_areas$fill),
                                          unique(shaded_areas$label)),
                        name = "Period") +
      
      # Customise labels and titles
      labs(title = paste("Overall Performance of", var_name, "during 20 August 2019 to 20 August 2020"),
           subtitle = paste("( Group by", label_name, ")"),
           x = "Date", 
           y = paste("Standardised", var_name)) +
      
      # Customise the x-axis and y-axis
      scale_x_date(date_labels = "%Y-%m-%d", 
                   date_breaks = "1 month") +
      theme_minimal() +
      theme(axis.text.x = element_text(angle = 55, 
                                       hjust = 1)) +
      scale_y_continuous(expand = expansion(mult = c(0.05, 0.1)))
    
    # Save the plot
    save_plot(p, paste0(dataset_name, "_", var_name))
    
    return(p)
  }
  
  dataset_name <- deparse(substitute(data))
  
  plots <- lapply(var_names, function(var_name) {
    create_plot(var_name, dataset_name)
  })
  
  # Print plots
  for (plot in plots) {
    print(plot)
  }
}

# Define the variables and their corresponding names
variables_scaled <- c('scaled_RET', 'scaled_volatility', 'scaled_sharpe_ratio', 'scaled_spread', 'scaled_dollar_vol', 'scaled_numtrd', 'scaled_turnover_ratio')
var_names <- c('Returns', 'Volatility', 'Sharpe Ratio', 'Bid-Ask Spread', 'Dollar Volume', 'Number of Trade', 'Turnover Ratio')

# Define the dates and colors for shaded areas
periods_group <- list(
  list(start = "2019-10-14", end = "2019-11-20", color = "#937613", label = "Non-COVID Period\n(2019-10-14 to 2019-11-20)"),
  list(start = "2020-02-14", end = "2020-03-20", color = "black", label = "COVID Period\n(2020-02-14 to 2020-03-20)")
  )

# Overall Performance
time_series_plot_overall(task1_overall, variables_scaled, var_names, periods_group)

# By security type
time_series_plot_group(task1_type, variables_scaled, var_names, "Security Types", periods_group, 'security_type', color_map = "Dark2")

# By industry sector
time_series_plot_group(task1_sector, variables_scaled, var_names,"Industry Sectors", periods_group, 'SECTOR', color_map = "Paired")

# By Company size
time_series_plot_group(task1_size, variables_scaled, var_names, "Company sizes", periods_group, 'company_size', color_map = "Set1")

market_risk <- function(data, variables, var_names, periods) {
  
  # Convert data to long format
  data_long <- data %>%
    pivot_longer(cols = all_of(variables), 
                 names_to = "variable", 
                 values_to = "value") %>%
    mutate(variable = factor(variable, 
                             levels = variables, 
                             labels = var_names)
           )
  
  # Create shaded areas data frame
  shaded_areas <- do.call(rbind, 
                          lapply(periods, function(p) {
                            data.frame(xmin = as.Date(p$start), 
                                       xmax = as.Date(p$end), 
                                       ymin = -Inf, 
                                       ymax = Inf, 
                                       fill = p$color, 
                                       label = p$label)
                            })
                          )
  
  # Define color mapping for each variable
  color_mapping <- c(
    "Returns" = "#f39906",
    "Volatility" = "#6402bb",
    "Sharpe Ratio" = "#32cd4c"
    )
  
# Define line type mapping
  linetype_mapping <- c(
    "Returns" = "solid",
    "Volatility" = "solid",
    "Sharpe Ratio" = "solid"
    )
  
  # Create the plot
  ggplot() +
    
    # Plot primary variables with unique colors and line types
    geom_line(data = data_long, 
              aes(x = date, 
                  y = value, 
                  color = variable, 
                  linetype = variable)
              ) +
    
    # Add shaded areas
    geom_rect(data = shaded_areas, 
              aes(xmin = xmin, 
                  xmax = xmax, 
                  ymin = ymin, 
                  ymax = ymax, 
                  fill = label),
              alpha = 0.2, 
              inherit.aes = FALSE) +
    
    # Define color and line type scales
    scale_color_manual(name = "Metrics",
                       values = color_mapping) +
    
    scale_linetype_manual(name = "Metrics", 
                          values = linetype_mapping) +
    
    # Define fill scales for shaded areas
    scale_fill_manual(values = setNames(unique(shaded_areas$fill),
                                        unique(shaded_areas$label)),
                      name = "Period") +
    
    # Define axis labels and limits
    labs(title = "Market Performance and Risk Metrics: August 20, 2019 to August 20, 2020", 
         x = "Date", 
         y = "Standardised Values of Performance Metrics") +
    
    scale_x_date(date_labels = "%Y-%m-%d", date_breaks = "1 month") +
    
    scale_y_continuous(limits = c(-5.5, 5.5), breaks = seq(-5.0, 5.0, by = 2)) +
    
    theme_minimal() +
    
    theme(axis.text.x = element_text(angle = 55, hjust = 1))
}

# Define variables
variables_risk <- c('scaled_RET', 'scaled_volatility', 'scaled_sharpe_ratio')
var_names_risk <- c('Returns', 'Volatility', 'Sharpe Ratio')

# Call function to create plot
performance_risk <- market_risk(task1_overall, variables_risk, var_names_risk, periods_group)

print(performance_risk)

# Save the graph
save_plot(performance_risk, "Time Series Plot for Market Performance and Risk Metrics")

market_liquidity <- function(data, variables, var_names, periods) {
  
  # Convert data to long format
  data_long <- data %>%
    pivot_longer(cols = all_of(variables), 
                 names_to = "variable", 
                 values_to = "value") %>%
    mutate(variable = factor(variable, 
                             levels = variables, 
                             labels = var_names)
           )

  
  # Create shaded areas data frame
  shaded_areas <- do.call(rbind, 
                          lapply(periods, function(p) {
                            data.frame(xmin = as.Date(p$start), 
                                       xmax = as.Date(p$end), 
                                       ymin = -Inf, 
                                       ymax = Inf, 
                                       fill = p$color, 
                                       label = p$label)
                            })
                          )
  
  # Define color mapping for each variable
  color_mapping <- c(
    "Bid-Ask Spread" = "#f39906",
    "Dollar Volume" = "#6402bb",
    "Number of Trade" = "#e85697",
    "Turnover Ratio" = "#32cd4c",
    "Volatility" = "black"
    )
  
# Define line type mapping (solid for all except Max Drawdown)
  linetype_mapping <- c(
    "Bid-Ask Spread" = "solid",
    "Dollar Volume" = "solid",
    "Turnover Ratio" = "solid",
    "Number of Trade" = "solid",
    "Volatility" = "dashed"
    )
  
  # Create the plot
  ggplot() +
    
    # Plot primary variables with unique colors and line types
    geom_line(data = data_long, 
              aes(x = date, 
                  y = value, 
                  color = variable, 
                  linetype = variable)
              ) +
    
    # Add shaded areas
    geom_rect(data = shaded_areas, 
              aes(xmin = xmin, 
                  xmax = xmax, 
                  ymin = ymin, 
                  ymax = ymax, 
                  fill = label),
              alpha = 0.2, 
              inherit.aes = FALSE) +
    
    # Define color and line type scales
    scale_color_manual(name = "Metrics",
                       values = color_mapping) +
    
    scale_linetype_manual(name = "Metrics", 
                          values = linetype_mapping) +
    
    # Define fill scales for shaded areas
    scale_fill_manual(values = setNames(unique(shaded_areas$fill),
                                        unique(shaded_areas$label)),
                      name = "Period") +
    
    # Define axis labels and limits
    labs(title = "Market Liquidity and Trading Behavior Metrics: August 20, 2019 to August 20, 2020",
         x = "Date", 
         y = "Standardised Values of Performance Metrics") +
    
    scale_x_date(date_labels = "%Y-%m-%d", date_breaks = "1 month") +
    
    scale_y_continuous(limits = c(-3, 6), breaks = seq(-3, 6, by = 2)) +
    
    theme_minimal() +
    
    theme(axis.text.x = element_text(angle = 55, hjust = 1))
}

# Define variables
variables_liquidity <- c('scaled_volatility', 'scaled_spread', 'scaled_dollar_vol', 'scaled_numtrd', 'scaled_turnover_ratio')

var_names_liquidity <- c('Volatility', 'Bid-Ask Spread', 'Dollar Volume', 'Number of Trade', 'Turnover Ratio')

# Call function to create plot
liquidity_trading <- market_liquidity(task1_overall, variables_liquidity, var_names_liquidity, periods_group)
print(liquidity_trading)

# Save the graph
save_plot(liquidity_trading, "Time Series Plot for Market Liquidity and Trading Behavior Metrics")

task2 <- dataset1_df9

# Prepare dataset
noncovid_filtered <- task2 %>% 
  filter(date < "2019-12-14") %>% 
  select(PERMNO, date, RET) %>% 
  left_join(fama_french_factors, by = "date")

# Initialise an empty list to store beta values
beta_list <- list()

# Loop through each unique PERMNO
for (permno in unique(noncovid_filtered$PERMNO)) {
  # Filter the data for the current PERMNO
  FFF_beta_noncovid <- subset(noncovid_filtered, PERMNO == permno)
  
  # Run the regression to calculate betas
  model_beta <- lm(RET ~ Mkt.RF + SMB + HML, data = FFF_beta_noncovid)
  
  # Store the betas with the PERMNO
  beta_list[[permno]] <- coef(model_beta)
}

# Convert the list of betas to a data frame
# Convert list to data frame with PERMNO names as row names
FFF_beta_noncovid <- do.call(rbind, lapply(names(beta_list), function(x) {
  as.data.frame(t(beta_list[[x]]), row.names = x)
}))

# Reset row names to a column
FFF_beta_noncovid <- tibble::rownames_to_column(FFF_beta_noncovid, var = "PERMNO")

# Rename columns to reflect beta coefficients
FFF_beta_noncovid <- FFF_beta_noncovid %>%
  rename(
    Mkt_beta = Mkt.RF,
    SMB_beta = SMB,
    HML_beta = HML
  ) %>% 
  select(-"(Intercept)")

FFF_beta_noncovid$PERMNO <- as.character(FFF_beta_noncovid$PERMNO)

print(head(FFF_beta_noncovid, n = 5))

# Prepare dataset
covid_filtered <- task2 %>% 
  filter(date < "2020-02-14") %>% 
  select(PERMNO, date, RET) %>% 
  left_join(fama_french_factors, by = "date")

# Initialise an empty list to store beta values
beta_list <- list()

# Loop through each unique PERMNO
for (permno in unique(covid_filtered$PERMNO)) {
  # Filter the data for the current PERMNO
  FFF_beta_covid <- subset(covid_filtered, PERMNO == permno)
  
  # Run the regression to calculate betas
  model_beta <- lm(RET ~ Mkt.RF + SMB + HML, data = FFF_beta_covid)
  
  # Store the betas with the PERMNO
  beta_list[[permno]] <- coef(model_beta)
}

# Convert the list of betas to a data frame
# Convert list to data frame with PERMNO names as row names
FFF_beta_covid <- do.call(rbind, lapply(names(beta_list), function(x) {
  as.data.frame(t(beta_list[[x]]), row.names = x)
}))

# Reset row names to a column
FFF_beta_covid <- tibble::rownames_to_column(FFF_beta_covid, var = "PERMNO")

# Rename columns to reflect beta coefficients
FFF_beta_covid <- FFF_beta_covid %>%
  rename(
    Mkt_beta = Mkt.RF,
    SMB_beta = SMB,
    HML_beta = HML
  ) %>% 
  select(-"(Intercept)")

FFF_beta_covid$PERMNO <- as.character(FFF_beta_covid$PERMNO)

print(head(FFF_beta_covid, n = 5))

transform_dataset_for_task2 <- function(data, start_date_RET, end_date_RET, start_date_others, end_date_others, type_of_security, FFF_dataset) {
  epsilon <- 1e-6
  
  # Find the date with the most PERMNOs having the highest price
  highest_change <- data %>%
    filter(date >= as.Date(start_date_RET) & date <= as.Date(end_date_RET)) %>%
    group_by(PERMNO) %>%
    filter(PRC == max(PRC)) %>%
    ungroup() %>%
    group_by(date) %>%
    summarise(num_highest = n()) %>%
    arrange(desc(num_highest)) 
  
  highest_date <- highest_change$date[1]
  cat("Highest Date: ")
  print(highest_change$date[1])

  # Find the date with the most PERMNOs having the lowest price
  lowest_change <- data %>%
    filter(date >= as.Date(start_date_RET) & date <= as.Date(end_date_RET)) %>%
    group_by(PERMNO) %>%
    filter(PRC == min(PRC)) %>%
    ungroup() %>%
    group_by(date) %>%
    summarise(num_lowest = n()) %>%
    arrange(desc(num_lowest))
  
  lowest_date <- lowest_change$date[1]
  cat("\nLowest Date: ")
  print(lowest_change$date[1])
  
  # Calculate the return using the lowest_date's and highest_date's price
  data_RET <- data %>%
    filter(date %in% c(highest_date, lowest_date)) %>%
    select(PERMNO, date, PRC) %>%
    spread(key = date, value = PRC, fill = NA) %>%
    rename(price_highest = !!as.character(highest_date), 
           price_lowest = !!as.character(lowest_date)) %>%
    mutate(returns = (price_highest - price_lowest) / price_lowest) %>%
    select(PERMNO, returns)
  
  cat("\nNumber of PERMNO in the dataset that calculated returns: ", length(unique(data_RET$PERMNO)), "\n")
  
  # Filter the data for two periods 
  data_filtered <- data %>%
    filter(date >= as.Date(start_date_others) & date <= as.Date(end_date_others)) %>%
    arrange(PERMNO, date)

  # Filter securities type and calculate aggregated variables
  data_others <- data_filtered %>%
    filter(
      type_of_security == "both" | 
      (type_of_security == "ETFs" & SHRCD == 73) |
      (type_of_security == "stocks" & SHRCD == 11)
    ) %>%
    group_by(PERMNO) %>%
    summarise(
      SHRCD = first(SHRCD),
      TICKER = first(TICKER),
      SECTOR = first(SECTOR),
      company_size = first(company_size),
      VOL = sum(VOL, na.rm = TRUE),
      dollar_vol = sum(dollar_vol, na.rm = TRUE),
      bid_ask_spread = mean(bid_ask_spread, na.rm = TRUE),
      turnover_ratio = mean(turnover_ratio, na.rm = TRUE),
      volatility = mean(volatility, na.rm = TRUE),
      sharpe_ratio = mean(sharpe_ratio, na.rm = TRUE),
      market_share = sum(market_share, na.rm = TRUE),
      market_cap = sum(market_cap, na.rm = TRUE),
      abs_corr_sp = mean(abs_corr_sp, na.rm = TRUE),
      security_type = first(security_type),
      tracking_error = mean(tracking_error, na.rm = TRUE),
      log_vol = log(sum(VOL + epsilon, na.rm = TRUE)),
      log_dollar_vol= log(sum(dollar_vol + epsilon, na.rm = TRUE)),
      log_spread = log(mean(bid_ask_spread + epsilon, na.rm = TRUE)),
      log_turnover_ratio = log(mean(turnover_ratio + epsilon, na.rm = TRUE)),
      log_volatility = log(mean(volatility + epsilon, na.rm = TRUE)),
      log_market_share = log(sum(market_share + epsilon, na.rm = TRUE)),
      log_market_cap = log(sum(market_cap + epsilon, na.rm = TRUE)),
      log_tracking_error = log(mean(tracking_error + epsilon, na.rm = TRUE)),
      vol_cap_ratio = VOL / market_cap,
      log_vol_cap_ratio = log(vol_cap_ratio),
      .groups = "drop"
    ) %>%
    mutate(
      scaled_vol = scale(VOL),
      scaled_dollar_vol = scale(dollar_vol),
      scaled_spread = scale(bid_ask_spread),
      scaled_turnover_ratio = scale(turnover_ratio),
      scaled_volatility = scale(volatility),
      scaled_sharpe_ratio = scale(sharpe_ratio),
      scaled_market_share = scale(market_share),
      scaled_market_cap = scale(market_cap),
      scaled_abs_corr_sp = scale(abs_corr_sp),
      scaled_tracking_error = scale(tracking_error)
    ) 
  
  # Check for duplicates
  duplicates <- data_others %>%
    group_by(PERMNO) %>%
    filter(n() > 1) %>%
    summarise(count = n())
  
  # Debugging
  if(nrow(duplicates) > 0) {
    warning("There are duplicates: ", paste(duplicates$PERMNO, collapse = ", "))
  }
  
  cat("Number of PERMNO that are in the dataset that aggregate x variables: ", length(unique(data_others$PERMNO)), "\n")
  
  cat("Number of PERMNO that are in the both datasets: ", length(unique(data_RET$PERMNO[data_RET$PERMNO %in% data_others$PERMNO])), "\n")
  
  # Join the returns values with other variables
  final_data <- data_RET %>%
    filter(PERMNO %in% data_others$PERMNO) %>%
    left_join(data_others, by = "PERMNO") %>%
    ungroup() %>%  
    left_join(FFF_dataset, by = "PERMNO") %>% 
    filter(returns != 0)
  
  cat("Number of PERMNO that are in the final dataset: ", length(unique(final_data$PERMNO)), "\n")
  
  return(final_data)
}

#Function to remove outlier
remove_outliers_all <- function(data, variables) {
  # Initialise a logical vector to keep track of rows to retain
  keep_rows <- rep(TRUE, nrow(data))
  
  # Loop over each variable
  for (variable in variables) {
    # Compute outlier bounds
    bounds <- data %>%
      summarise(
        Q1 = quantile(.data[[variable]], 0.25, na.rm = TRUE),
        Q3 = quantile(.data[[variable]], 0.75, na.rm = TRUE),
        IQR = Q3 - Q1,
        lower_bound = Q1 - 1.5 * IQR,
        upper_bound = Q3 + 1.5 * IQR
      )
    
    # Apply the filtering based on computed bounds
    keep_rows <- keep_rows & (data[[variable]] >= bounds$lower_bound & data[[variable]] <= bounds$upper_bound)
  }
  
  # Return the cleaned data
  return(data[keep_rows, ])
}

#Find the correlation between retruns and each potential explanatory variables
cor_p1_m1 <- cor(task2_reg_p1_df1 %>% 
                   select(-c(PERMNO, SHRCD, TICKER, SECTOR, company_size, security_type)), 
                 use = "pairwise.complete.obs")

# Correlation over absolute 0.7 = high risk of multicollinearity
high_corr_p1_m1 <- which(abs(cor_p1_m1) > 0.7, arr.ind = TRUE)
high_corr_pairs_p1_m1 <- data.frame(
  Feature1 = rownames(cor_p1_m1)[high_corr_p1_m1[,1]],
  Feature2 = colnames(cor_p1_m1)[high_corr_p1_m1[,2]],
  Correlation = cor_p1_m1[high_corr_p1_m1]
)

high_corr_pairs_p1_m1 <- high_corr_pairs_p1_m1 %>% 
  filter(Correlation < 1)

# Inspect the variables with multicollinearity
head(high_corr_pairs_p1_m1)

# Inspect all correlations
head(melt(cor_p1_m1))

# Splitting the dataset into training and testing sets
set.seed(696)  

# Non-COVID Period
X_p1_m1 <- task2_reg_p1_df1 %>% 
  select(Mkt_beta, SMB_beta, HML_beta, log_turnover_ratio, scaled_spread, log_volatility, log_market_share, vol_cap_ratio, company_sizeSmall, company_sizeLarge, SECTORMaterials, SECTORIndustrials, SECTORFinancials, SECTORInformation_Technology, SECTORConsumer_Discretionary, SECTORHealth_Care, SECTOREnergy, SECTORUtilities, SECTORReal_Estate, SECTORCommunication_Services, SECTORUnknown)

y_p1_m1 <- task2_reg_p1_df1$returns

splitIndex_p1_m1 <- createDataPartition(y_p1_m1, 
                                        p = 0.8, 
                                        list = FALSE)
X_train_data_p1_m1 <- X_p1_m1[splitIndex_p1_m1, ]
X_test_data_p1_m1 <- X_p1_m1[-splitIndex_p1_m1, ]
y_train_data_p1_m1 <- y_p1_m1[splitIndex_p1_m1]
y_test_data_p1_m1 <- y_p1_m1[-splitIndex_p1_m1]

train_data_p1_m1 <- cbind(X_train_data_p1_m1, returns = y_train_data_p1_m1)

# Build Model
model_1_p1 <- lm(returns ~ ., data = train_data_p1_m1)

# Prediction for training sample
pred_train_p1_m1 <- predict(model_1_p1, newdata = X_train_data_p1_m1)

metrics_train_p1_m1 <- postResample(pred = pred_train_p1_m1, 
                                    obs = train_data_p1_m1$returns)
metrics_table_train_p1_m1 <- data.frame(
  Metric = c('MSE', 'RMSE', 'MAE', 'R-squared'),
  Value = c(metrics_train_p1_m1["RMSE"]^2,
            metrics_train_p1_m1["RMSE"], 
            metrics_train_p1_m1["MAE"], 
            metrics_train_p1_m1["Rsquared"])
)

# Prediction for testing sample
pred_test_p1_m1 <- predict(model_1_p1, newdata = X_test_data_p1_m1)
metrics_test_p1_m1 <- postResample(pred = pred_test_p1_m1, 
                                   obs = y_test_data_p1_m1)
metrics_table_test_p1_m1 <- data.frame(
  Metric = c('MSE', 'RMSE', 'MAE', 'R-squared'),
  Value = c(metrics_test_p1_m1["RMSE"]^2
, metrics_test_p1_m1["RMSE"], metrics_test_p1_m1["MAE"], metrics_test_p1_m1["Rsquared"])
)

# Scatterplot of Prediction vs Actual
plot_data_p1_m1 <- data.frame(
  Predicted = pred_test_p1_m1,
  Actual = y_test_data_p1_m1
)

ggplot(data = plot_data_p1_m1,
       aes(x = Predicted,
           y = Actual)) +
  
  geom_point() +
  
  geom_smooth(method = "lm",
              se = FALSE,
              color = "blue",
              aes(linetype = "Regression Line")) +
  
  scale_linetype_manual(name = NULL,
                        values = "solid",
                        labels = "Regression Line") +
  
  theme_minimal() +
  
  theme(legend.position = "bottom") +
  
  labs(title = "Evaluating Model 1's Performance: Predicted vs. Actual Returns in the Non-COVID Period",
       x = "Predicted returns",
       y = "Actual returns")

# Extract coefficients
coef_m1 <- coef(model_1_p1)

importance_m1 <- data.frame(
  Factors = names(coef_m1),
  Coefficient = coef_m1
)

importance_m1 <- importance_m1[importance_m1$Factor != "(Intercept)", ]

# Plot with positive and negative impacts
plot_m1_p1_factors <- ggplot(importance_m1, 
                             aes(x = reorder(Factors, Coefficient), 
                                 y = Coefficient, fill = Coefficient > 0)) +
  
  geom_bar(stat = "identity") +
  
  geom_text(aes(label = round(Coefficient, 3)),
            vjust = ifelse(importance_m1$Coefficient > 0, 0.5, 0.3),
            color = "black") +
  
  coord_flip() +
  
  labs(title = "Impact of factors on Returns for Model 1",
       x = "Factors",
       y = "Coefficient") +
  
  theme_minimal() +
  
  scale_fill_manual(
    values = c("TRUE" = "steelblue", "FALSE" = "coral"),
    labels = c("TRUE" = "Positive Impact", 
               "FALSE" = "Negative Impact"),
    name = ""
  ) +
  
  theme(
    axis.title.y = element_text(color = 'black'),
    axis.text.y = element_text(color = 'black'),
    legend.position = 'bottom'
  ) 

print(plot_m1_p1_factors)

# Apply the model particularly by industry
pred_sector_p1_m1  <- cbind(X_test_data_p1_m1, SECTOR = task2_reg_p1_df1$SECTOR[-splitIndex_p1_m1])

data_sector_p1_m1 <- data.frame(
  Predicted = pred_test_p1_m1,
  Actual = y_test_data_p1_m1,
  SECTOR = pred_sector_p1_m1$SECTOR
)

data_sector_p1_m1$SECTOR <- as.character(data_sector_p1_m1$SECTOR)

# Calculate evaluation metrics to see the performance of model apply on SECTOR
metrics_table_sector_p1_m1 <- data_sector_p1_m1 %>% 
  group_by(SECTOR) %>% 
  summarise(
    MSE = mean((Actual - Predicted)^2),
    RMSE = sqrt(MSE),
    MAE = mean(abs(Actual - Predicted)),
    residuals = sum((Actual - Predicted)^2),
    total = sum((Actual - mean(Actual))^2),
    R_Squared = 1 - (residuals / total)
  ) %>% 
  select(-c(residuals, total)) 

print(metrics_table_sector_p1_m1)

# Plot the prediction vs actual by SECTOR
plot_m1_p1_sector <- ggplot(data = data_sector_p1_m1, 
                     aes(x = Predicted, 
                         y = Actual,
                         color = SECTOR)) +
  
  geom_point(size = 2) +
  
  geom_smooth(method = "lm", 
              se = FALSE, 
              color = "blue", 
              aes(linetype = "Regression Line")) + 
  
  scale_linetype_manual(name = NULL, 
                        values = "solid", 
                        labels = "Regression Line") + 
  
  theme_minimal() +
  
  theme(legend.position = "bottom",
        plot.title = element_text(size = 14, face = "bold"),
        plot.subtitle = element_text(size = 12),
        strip.text = element_text(size = 12, face = "bold")) + 
  
  labs(title = "Evaluating Model 1's Performance: Predicted vs. Actual Returns",
       subtitle = "During the Non-COVID Period (Grouped by Industry Sector)",  
       x = "Predicted returns", 
       y = "Actual returns",
       color = "Industry Sector")


print(plot_m1_p1_sector)

# Apply the model particularly by security type
pred_type_p1_m1  <- cbind(X_test_data_p1_m1, security_type = task2_reg_p1_df1$security_type[-splitIndex_p1_m1])

data_type_p1_m1 <- data.frame(
  Predicted = pred_test_p1_m1,
  Actual = y_test_data_p1_m1,
  security_type = pred_type_p1_m1$security_type
)

# Calculate evaluation metrics to see the performance of model apply on security type
metrics_table_type_p1_m1 <- data_type_p1_m1 %>% 
  group_by(security_type) %>% 
  summarise(
    MSE = mean((Actual - Predicted)^2),
    RMSE = sqrt(MSE),
    MAE = mean(abs(Actual - Predicted)),
    residuals = sum((Actual - Predicted)^2),
    total = sum((Actual - mean(Actual))^2),
    R_Squared = 1 - (residuals / total)
  ) %>% 
  select(-c(residuals, total))

print(metrics_table_type_p1_m1)

# Plot the prediction vs actual by security type
plot_m1_p1_type <- ggplot(data = data_type_p1_m1, 
                     aes(x = Predicted, 
                         y = Actual,
                         color = security_type)) +
  geom_point(size = 2) +
  geom_smooth(method = "lm", 
              se = FALSE, 
              color = "blue", 
              aes(linetype = "Regression Line")) + 
  scale_linetype_manual(name = NULL, 
                        values = "solid", 
                        labels = "Regression Line") + 
  theme_minimal() +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 14, face = "bold"),
        plot.subtitle = element_text(size = 12),
        strip.text = element_text(size = 12, face = "bold")) +  
  labs(title = "Evaluating Model 1's Performance: Predicted vs. Actual Returns",
       subtitle = "During the Non-COVID Period (Grouped by Security Type)",  
       x = "Predicted returns", 
       y = "Actual returns",
       color = "Security Type")


print(plot_m1_p1_type)

# Apply the model particularly by company size
pred_size_p1_m1  <- cbind(X_test_data_p1_m1, company_size = task2_reg_p1_df1$company_size[-splitIndex_p1_m1])

data_size_p1_m1 <- data.frame(
  Predicted = pred_test_p1_m1,
  Actual = y_test_data_p1_m1,
  company_size = pred_size_p1_m1$company_size
)

# Calculate evaluation metrics to see the performance of model apply on company size
metrics_table_size_p1_m1 <- data_size_p1_m1 %>% 
  group_by(company_size) %>% 
  summarise(
    MSE = mean((Actual - Predicted)^2),
    RMSE = sqrt(MSE),
    MAE = mean(abs(Actual - Predicted)),
    residuals = sum((Actual - Predicted)^2),
    total = sum((Actual - mean(Actual))^2),
    R_Squared = 1 - (residuals / total)
  ) %>% 
  select(-c(residuals, total))

print(metrics_table_size_p1_m1)

# Plot the prediction vs actual by company size
plot_m1_p1_size <- ggplot(data = data_size_p1_m1, 
                     aes(x = Predicted, 
                         y = Actual,
                         color = company_size)) +
  geom_point(size = 2, alpha = 0.7) +
  geom_smooth(method = "lm", 
              se = FALSE, 
              color = "blue", 
              aes(linetype = "Regression Line")) + 
  scale_linetype_manual(name = NULL, 
                        values = "solid", 
                        labels = "Regression Line") + 
  theme_minimal() +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 14, face = "bold"),
        plot.subtitle = element_text(size = 12),
        strip.text = element_text(size = 12, face = "bold")) +  
  labs(title = "Evaluating Model 1's Performance: Predicted vs. Actual Returns",
       subtitle = "During the Non-COVID Period (Grouped by Company Size)",  
       x = "Predicted returns", 
       y = "Actual returns",
       color = "Company Size")


print(plot_m1_p1_size)

# Find the date with the most PERMNOs having the highest price change
highest_change_covid <- task2 %>%
  filter(date >= "2020-02-14" & date <= "2020-03-20") %>%
  group_by(PERMNO) %>%
  filter(PRC == max(PRC)) %>%
  ungroup() %>%
  group_by(date) %>%
  summarise(num_highest = n()) %>%
  arrange(desc(num_highest))

# Find the date with the most PERMNOs having the lowest price change
lowest_change_covid <- task2 %>%
  filter(date >= "2020-02-14" & date <= "2020-03-20") %>%
  group_by(PERMNO) %>%
  filter(PRC == min(PRC)) %>%
  ungroup() %>%
  group_by(date) %>%
  summarise(num_lowest = n()) %>%
  arrange(desc(num_lowest))

# Inspect the results
returns_covid <- task2 %>%
  filter(date == highest_change_covid$date[1] | date == lowest_change_covid$date[1]) %>% 
  group_by(PERMNO) %>% 
  summarise(max_change = PRC[date == highest_change_covid$date[1]],
    min_change = PRC[date == lowest_change_covid$date[1]],
    returns_covid = (max_change - min_change) / min_change) %>% 
    select(PERMNO, returns_covid)

comparison_covid <- returns_covid  %>%
  left_join(task2_reg_p2_df1, by = "PERMNO") %>%
  mutate(difference_covid = returns_covid != returns) %>%
  filter(difference_covid) 

#Find the correlation between retruns and each potential explanatory variables
cor_p2_m1 <- cor(task2_reg_p2_df1%>% select(-c(PERMNO, SHRCD, TICKER, SECTOR, company_size, security_type)), use = "pairwise.complete.obs")

# Correlation over absolute 0.7 = high risk of multicollinearity
high_corr_p2_m1 <- which(abs(cor_p2_m1) > 0.7, arr.ind = TRUE)
high_corr_pairs_p2_m1 <- data.frame(
  Feature1 = rownames(cor_p2_m1)[high_corr_p2_m1[,1]],
  Feature2 = colnames(cor_p2_m1)[high_corr_p2_m1[,2]],
  Correlation = cor_p2_m1[high_corr_p2_m1]
)
high_corr_pairs_p2_m1 <- high_corr_pairs_p2_m1 %>% filter(Correlation < 1)

# Inspect the variables with multicollinearity
head(high_corr_pairs_p2_m1)

# Inspect all correlations
head(melt(cor_p2_m1))

# Splitting the dataset into training and testing sets
set.seed(781)  

# COVID Period
X_p2_m2 <- task2_reg_p2_df1 %>% 
  select(Mkt_beta, SMB_beta, HML_beta, log_volatility, log_spread, scaled_market_share, scaled_abs_corr_sp, scaled_dollar_vol, SECTORMaterials, SECTORIndustrials, SECTORFinancials, SECTORInformation_Technology, SECTORConsumer_Discretionary, SECTORHealth_Care, SECTOREnergy, SECTORUtilities, SECTORReal_Estate, SECTORCommunication_Services, SECTORUnknown)

y_p2_m2 <- task2_reg_p2_df1$returns

splitIndex_p2_m2 <- createDataPartition(y_p2_m2, 
                                        p = 0.8, 
                                        list = FALSE)
X_train_data_p2_m2 <- X_p2_m2[splitIndex_p2_m2, ]
X_test_data_p2_m2 <- X_p2_m2[-splitIndex_p2_m2, ]
y_train_data_p2_m2 <- y_p2_m2[splitIndex_p2_m2]
y_test_data_p2_m2 <- y_p2_m2[-splitIndex_p2_m2]

train_data_p2_m2 <- cbind(X_train_data_p2_m2, returns = y_train_data_p2_m2)

# Build Model
model_2_p2 <- lm(returns ~ ., data = train_data_p2_m2)

# Prediction for training sample - Covid
pred_train_p2_m2 <- predict(model_2_p2, newdata = X_train_data_p2_m2)

metrics_train_p2_m2 <- postResample(pred = pred_train_p2_m2, 
                                    obs = train_data_p2_m2$returns)
metrics_table_train_p2_m2 <- data.frame(
  Metric = c('MSE', 'RMSE', 'MAE', 'R-squared'),
  Value = c(metrics_train_p2_m2["RMSE"]^2,
            metrics_train_p2_m2["RMSE"], 
            metrics_train_p2_m2["MAE"], 
            metrics_train_p2_m2["Rsquared"])
)

# Prediction for testing sample - covid
pred_test_p2_m2 <- predict(model_2_p2, newdata = X_test_data_p2_m2)
metrics_test_p2_m2 <- postResample(pred = pred_test_p2_m2, 
                                   obs = y_test_data_p2_m2)
metrics_table_test_p2_m2 <- data.frame(
  Metric = c('MSE', 'RMSE', 'MAE', 'R-squared'),
  Value = c(metrics_test_p2_m2["RMSE"]^2
, metrics_test_p2_m2["RMSE"], metrics_test_p2_m2["MAE"], metrics_test_p2_m2["Rsquared"])
)

# Scatterplot of Prediction vs Actual
plot_data_p2_m2 <- data.frame(
  Predicted = pred_test_p2_m2,
  Actual = y_test_data_p2_m2
)

plot_m2_p2 <- ggplot(data = plot_data_p2_m2,
                     aes(x = Predicted,
                         y = Actual)) +
  geom_point() +
  geom_smooth(method = "lm",
              se = FALSE,
              color = "blue",
              aes(linetype = "Regression Line")) +
  scale_linetype_manual(name = NULL,
                        values = "solid",
                        labels = "Regression Line") +
  theme_minimal() +
  theme(legend.position = "bottom") +
  labs(title = "Evaluating Model 2's Performance: Predicted vs. Actual Returns in the Non-covid Period",
       x = "Predicted returns",
       y = "Actual returns")

print(plot_m2_p2)

# Extract coefficients
coef_m2 <- coef(model_2_p2)

importance_m2 <- data.frame(
  Factors = names(coef_m2),
  Coefficient = coef_m2
)
importance_m2 <- importance_m2[importance_m2$Factor != "(Intercept)", ]
# Plot with positive and negative impacts
plot_m2_p2_factors <- ggplot(importance_m2, aes(x = reorder(Factors, Coefficient), y = Coefficient, fill = Coefficient > 0)) +
  geom_bar(stat = "identity") +
  geom_text(aes(label = round(Coefficient, 3)),
            vjust = ifelse(importance_m2$Coefficient > 0, 0.5, 0.3),
            color = "black") +
  coord_flip() +
  labs(title = "Impact of factors on returns for Model 2",
       x = "Factors",
       y = "Coefficient") +
  theme_minimal() +
  scale_fill_manual(
    values = c("TRUE" = "steelblue", "FALSE" = "coral"),
    labels = c("TRUE" = "Positive Impact", 
               "FALSE" = "Negative Impact"),
    name = ""
  ) +
  theme(
    axis.title.y = element_text(color = 'black'),
    axis.text.y = element_text(color = 'black'),
    legend.position = 'bottom'
  ) 

pred_sector_p2_m2  <- cbind(X_test_data_p2_m2, SECTOR = task2_reg_p2_df1$SECTOR[-splitIndex_p2_m2])

data_sector_p2_m2 <- data.frame(
  Predicted = pred_test_p2_m2,
  Actual = y_test_data_p2_m2,
  SECTOR = pred_sector_p2_m2$SECTOR
)

data_sector_p2_m2$SECTOR <- as.character(data_sector_p2_m2$SECTOR)

# Calculating performance metrics by sector
metrics_table_sector_p2_m2 <- data_sector_p2_m2 %>% 
  group_by(SECTOR) %>% 
  summarise(
    MSE = mean((Actual - Predicted)^2),
    RMSE = sqrt(MSE),
    MAE = mean(abs(Actual - Predicted)),
    residuals = sum((Actual - Predicted)^2),
    total = sum((Actual - mean(Actual))^2),
    R_Squared = 1 - (residuals / total)
  ) %>% 
  select(-c(residuals, total))

print(metrics_table_sector_p2_m2)

plot_m2_p2_sector <- ggplot(data = data_sector_p2_m2, 
                     aes(x = Predicted, 
                         y = Actual,
                         color = SECTOR)) +
  geom_point(size = 2) +
  geom_smooth(method = "lm", 
              se = FALSE, 
              color = "blue", 
              aes(linetype = "Regression Line")) + 
  scale_linetype_manual(name = NULL, 
                        values = "solid", 
                        labels = "Regression Line") + 
  theme_minimal() +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 14, face = "bold"),
        plot.subtitle = element_text(size = 12),
        strip.text = element_text(size = 12, face = "bold")) +  
  labs(title = "Evaluating Model 2's Performance: Predicted vs. Actual Returns",
       subtitle = "During the COVID Period (Grouped by Industry Sector)",  
       x = "Predicted returns", 
       y = "Actual returns",
       color = "Industry Sector")


print(plot_m2_p2_sector)